from pypcd import pypcd
from point_to_bev import *
import os
import json
from scipy.spatial.transform import Rotation as R
import numpy as np
import cv2


def rotation_matrix_x(theta):
    """绕x轴的旋转矩阵"""
    return np.array([
        [1, 0, 0],
        [0, np.cos(theta), -np.sin(theta)],
        [0, np.sin(theta), np.cos(theta)]
    ])

def rotation_matrix_y(theta):
    """绕y轴的旋转矩阵"""
    return np.array([
        [np.cos(theta), 0, np.sin(theta)],
        [0, 1, 0],
        [-np.sin(theta), 0, np.cos(theta)]
    ])

def rotation_matrix_z(theta):
    """绕z轴的旋转矩阵"""
    return np.array([
        [np.cos(theta), -np.sin(theta), 0],
        [np.sin(theta), np.cos(theta), 0],
        [0, 0, 1]
    ])

def calculate_vertices(x, y, z, l, w, h, yaw):
    # 长方体的半尺寸
    half_l = l / 2
    half_w = w / 2
    half_h = h / 2

    # 未旋转时的顶点坐标
    vertices = np.array([
        [x - half_l, y - half_w, z + half_h],
        [x + half_l, y - half_w, z + half_h],
        [x + half_l, y + half_w, z + half_h],
        [x - half_l, y + half_w, z + half_h],
        [x - half_l, y - half_w, z - half_h],
        [x + half_l, y - half_w, z - half_h],
        [x + half_l, y + half_w, z - half_h],
        [x - half_l, y + half_w, z - half_h]
    ])

    # 绕z轴旋转的变换矩阵
    # rx = rotation_matrix_x(roll)
    # ry = rotation_matrix_y(pitch)
    rz = rotation_matrix_z(yaw)

    # 旋转顶点坐标
    # rotated_vertices = np.dot(vertices - np.array([x, y, z]), rx) + np.array([x, y, z])
    # rotated_vertices = np.dot(vertices - np.array([x, y, z]), ry) + np.array([x, y, z])
    rotated_vertices = np.dot(vertices - np.array([x, y, z]), rz) + np.array([x, y, z])

    return rotated_vertices